#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
列删除工具
版本: 1.0.0
作者: 
功能: 从Excel/CSV文件中删除指定列，根据列名列表文件批量删除不需要的列
"""

import argparse
import pandas as pd
from pathlib import Path
from typing import List, Optional
import sys

class ColumnRemover:
    def __init__(self, input_file: str, header_file: str, output_file: str, 
                 output_format: Optional[str] = None):
        self.input_file = Path(input_file)
        self.header_file = Path(header_file)
        self.output_file = Path(output_file)
        self.output_format = output_format
        self.data = None
        self.columns_to_remove = []
        
    def load_columns_to_remove(self) -> bool:
        """从文件加载需要删除的列名列表"""
        try:
            with open(self.header_file, 'r', encoding='utf-8') as f:
                # 读取文件中的每一行，去除空白字符
                self.columns_to_remove = [line.strip() for line in f if line.strip()]
            
            if not self.columns_to_remove:
                print(f"警告: 列名文件 {self.header_file} 为空")
                return False
                
            print(f"从文件加载了 {len(self.columns_to_remove)} 个需要删除的列名")
            print(f"需要删除的列: {self.columns_to_remove}")
            return True
            
        except Exception as e:
            print(f"加载列名文件失败: {e}")
            return False
    
    def load_data(self) -> bool:
        """加载数据文件"""
        try:
            file_extension = self.input_file.suffix.lower()
            
            if file_extension == '.xlsx':
                self.data = pd.read_excel(self.input_file)
            elif file_extension == '.xls':
                self.data = pd.read_excel(self.input_file)
            elif file_extension == '.csv':
                self.data = pd.read_csv(self.input_file)
            elif file_extension == '.tsv':
                self.data = pd.read_csv(self.input_file, sep='\t')
            else:
                raise ValueError(f"不支持的文件格式: {file_extension}")
                
            print(f"成功加载数据: {len(self.data)} 行, {len(self.data.columns)} 列")
            print(f"原始列名: {list(self.data.columns)}")
            return True
            
        except Exception as e:
            print(f"加载数据失败: {e}")
            return False
    
    def validate_columns(self) -> bool:
        """验证需要删除的列名是否存在"""
        if not self.columns_to_remove:
            print("错误: 没有需要删除的列")
            return False
            
        missing_columns = []
        for col in self.columns_to_remove:
            if col not in self.data.columns:
                missing_columns.append(col)
        
        if missing_columns:
            print(f"警告: 以下列不存在于数据文件中: {missing_columns}")
            print(f"可用列名: {list(self.data.columns)}")
            print(f"将删除存在的列: {[col for col in self.columns_to_remove if col in self.data.columns]}")
        
        return True
    
    def remove_columns(self) -> pd.DataFrame:
        """删除指定列"""
        try:
            # 获取实际存在的需要删除的列
            existing_columns_to_remove = [col for col in self.columns_to_remove if col in self.data.columns]
            
            if not existing_columns_to_remove:
                print("警告: 没有需要删除的列存在于数据文件中")
                return self.data.copy()
            
            # 删除列
            cleaned_data = self.data.drop(columns=existing_columns_to_remove)
            print(f"成功删除 {len(existing_columns_to_remove)} 列: {existing_columns_to_remove}")
            print(f"剩余 {len(cleaned_data.columns)} 列: {list(cleaned_data.columns)}")
            
            return cleaned_data
            
        except Exception as e:
            print(f"删除列失败: {e}")
            return None
    
    def save_result(self, cleaned_data: pd.DataFrame) -> bool:
        """保存清理后的数据"""
        try:
            # 确保输出目录存在
            self.output_file.parent.mkdir(parents=True, exist_ok=True)
            
            # 确定输出格式
            if self.output_format:
                output_extension = f".{self.output_format.lower()}"
            else:
                # 如果未指定输出格式，则使用输入文件的格式
                output_extension = self.output_file.suffix.lower()
                if not output_extension:
                    # 如果输出文件没有扩展名，则使用输入文件的扩展名
                    output_extension = self.input_file.suffix.lower()
            
            # 如果仍然没有确定格式，则默认使用CSV
            if not output_extension:
                output_extension = '.csv'
            
            print(f"保存文件格式: {output_extension}")
            
            if output_extension == '.xlsx':
                cleaned_data.to_excel(self.output_file, index=False)
            elif output_extension == '.xls':
                cleaned_data.to_excel(self.output_file, index=False)
            elif output_extension == '.csv':
                cleaned_data.to_csv(self.output_file, index=False)
            elif output_extension == '.tsv':
                cleaned_data.to_csv(self.output_file, sep='\t', index=False)
            else:
                # 默认保存为CSV格式
                print(f"未知格式 {output_extension}，默认保存为CSV格式")
                cleaned_data.to_csv(self.output_file, index=False)
            
            print(f"结果已保存到: {self.output_file}")
            return True
            
        except Exception as e:
            print(f"保存失败: {e}")
            return False
    
    def process(self) -> bool:
        """执行完整的删除流程"""
        print("开始列删除...")
        print(f"输入文件: {self.input_file}")
        print(f"列名列表文件: {self.header_file}")
        print(f"输出文件: {self.output_file}")
        
        # 1. 加载需要删除的列名
        if not self.load_columns_to_remove():
            return False
        
        # 2. 加载数据
        if not self.load_data():
            return False
        
        # 3. 验证列名
        if not self.validate_columns():
            return False
        
        # 4. 删除列
        cleaned_data = self.remove_columns()
        if cleaned_data is None:
            return False
        
        # 5. 保存结果
        if not self.save_result(cleaned_data):
            return False
        
        print("\n✅ 列删除完成!")
        print(f"📁 输入文件: {self.input_file}")
        print(f"📋 列名列表文件: {self.header_file}")
        print(f"📄 输出文件: {self.output_file}")
        print(f"📊 删除列数: {len([col for col in self.columns_to_remove if col in self.data.columns])}")
        print(f"📈 剩余列数: {len(cleaned_data.columns)}")
        
        return True

def main():
    parser = argparse.ArgumentParser(
        description='列删除工具 - 从Excel/CSV文件中删除指定列',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 从CSV文件中删除指定列
  column-remover -i data.csv --header headers_to_remove.txt -o cleaned_data.csv
  
  # 从Excel文件中删除指定列
  column-remover -i data.xlsx --header headers_to_remove.txt -o cleaned_data.xlsx
  
  # 列名列表文件格式（每行一个列名）
  # headers_to_remove.txt:
  # ID
  # Name  
  # Address
  # Phone

支持的文件格式: Excel (.xlsx, .xls), CSV (.csv), TSV (.tsv)
        """
    )
    parser.add_argument('-i', '--input', required=True, help='输入文件路径 (支持Excel/CSV/TSV格式)')
    parser.add_argument('--header', required=True, help='包含需要删除列名的列表文件路径（每行一个列名）')
    parser.add_argument('-o', '--output', required=True, help='输出文件路径')
    parser.add_argument('--format', choices=['xlsx', 'xls', 'csv', 'tsv'], 
                       help='输出文件格式（如果不指定，则与输入文件格式相同）')
    
    args = parser.parse_args()
    
    remover = ColumnRemover(
        args.input,
        args.header,
        args.output,
        args.format
    )
    
    success = remover.process()
    
    if success:
        print("\n🎉 列删除工具执行成功!")
        return 0
    else:
        print("\n❌ 列删除工具执行失败!")
        return 1

if __name__ == '__main__':
    sys.exit(main())